#if UNITY_EDITOR
using System.Collections.Generic;
using UnityEditor;
using UnityEngine;
using UnityEngine.Rendering;

public class GpuPVSBakerWindow : EditorWindow
{
    // ===== 用户参数 =====
    Vector3Int   _grid = new(10,4,10);
    Vector3      _min  = new(-5,-1,-5);
    Vector3      _max  = new( 5, 3, 5);
    int          _res  = 256;               // 每个面的渲染分辨率
    int          _samplesPerCell = 1;       // 每格采样点：1=中心；可改 3/14 etc.
    int          _hashSizePow2 = 4096;      // Compute 哈希表大小（2^n）

    List<Renderer> _occludees = new();      // 需要写位图的对象
    List<Renderer> _allRenderers = new();   // 场景全部 renderer（用于渲染ID）
    Dictionary<Renderer, int> _idOf = new();// 全场 renderer 的全局ID（0 保留为背景）
    Material _idMat;                        // ID 输出材质
    Dictionary<Renderer, Material[]> _backupMats = new();

    RenderTexture _rt;                      // ID 贴图
    ComputeShader _cs;                      // 去重核
    ComputeBuffer _hash, _unique;           // 去重缓存
    int _kernel;

    [MenuItem("Tools/PVS (GPU)")]
    static void Open() => GetWindow<GpuPVSBakerWindow>("PVS (GPU)");

    [MenuItem("Tools/PVS (GPU)/Create Sample Scene")]
    static void CreateSample()
    {
        var root = new GameObject("PVS_Sample");

        var ground = GameObject.CreatePrimitive(PrimitiveType.Plane);
        ground.name="Ground"; ground.transform.SetParent(root.transform);

        var wall = GameObject.CreatePrimitive(PrimitiveType.Cube);
        wall.name="Wall"; wall.transform.SetParent(root.transform);
        wall.transform.position = new Vector3(0,1,0);
        wall.transform.localScale= new Vector3(0.2f,2f,6f);

        var red = GameObject.CreatePrimitive(PrimitiveType.Sphere);
        red.name="Red"; red.transform.SetParent(root.transform);
        red.transform.position = new Vector3(-2,0.5f,0);

        var blue = GameObject.CreatePrimitive(PrimitiveType.Sphere);
        blue.name="Blue"; blue.transform.SetParent(root.transform);
        blue.transform.position = new Vector3(+2,0.5f,0);
        var m = new Material(Shader.Find("Universal Render Pipeline/Lit"));
        if (m) { m.SetColor("_BaseColor", new Color(0.2f,0.5f,1f,1f)); blue.GetComponent<Renderer>().sharedMaterial = m; }

        var camGO = new GameObject("Main Camera");
        var cam = camGO.AddComponent<Camera>();
        cam.tag="MainCamera";
        cam.transform.position = new Vector3(-4,2.5f,-4);
        cam.transform.rotation = Quaternion.Euler(20,35,0);
        cam.nearClipPlane=0.05f; cam.farClipPlane=60f;
        camGO.AddComponent<AudioListener>();

        Selection.activeObject = root;
    }

    void OnEnable()
    {
        _idMat = new Material(Shader.Find("Hidden/PVS/ObjectId"));
        _cs = AssetDatabase.LoadAssetAtPath<ComputeShader>("Assets/PVS/Shaders/PVS_CollectIDs.compute");
        if (!_cs) Debug.LogError("找不到 Compute：Assets/Shaders/PVS_CollectIDs.compute");
        _kernel = _cs ? _cs.FindKernel("Collect") : -1;
    }
    void OnDisable() { CleanupRTAndBuffers(); RestoreMaterials(); }

    void CleanupRTAndBuffers()
    {
        if (_rt) DestroyImmediate(_rt); _rt=null;
        _hash?.Dispose(); _hash=null;
        _unique?.Dispose(); _unique=null;
    }

    void OnGUI()
    {
        EditorGUILayout.LabelField("Grid Volume", EditorStyles.boldLabel);
        _min = EditorGUILayout.Vector3Field("Min", _min);
        _max = EditorGUILayout.Vector3Field("Max", _max);
        _grid = EditorGUILayout.Vector3IntField("Grid (X,Y,Z)", _grid);

        EditorGUILayout.Space(4);
        EditorGUILayout.LabelField("Bake Settings", EditorStyles.boldLabel);
        _res = EditorGUILayout.IntSlider("Face Resolution", _res, 64, 1024);
        _samplesPerCell = EditorGUILayout.IntSlider("Samples per Cell", _samplesPerCell, 1, 14);
        _hashSizePow2 = EditorGUILayout.IntPopup("Hash Size", _hashSizePow2,
            new[] { "4096", "8192", "16384" }, new[] { 4096, 8192, 16384 });

        if (GUILayout.Button("1) Collect Objects"))
        {
            _occludees.Clear(); _allRenderers.Clear(); _idOf.Clear();
            _allRenderers.AddRange(Object.FindObjectsOfType<Renderer>());
            // 过滤一些不需要的渲染器类型
            _allRenderers.RemoveAll(r => r is TrailRenderer || r is LineRenderer || r is ParticleSystemRenderer);

            // —— occludees（要出现在位图里的目标）：示例仅小球；实际可自定义筛选 —— //
            foreach (var r in _allRenderers)
                if (r && (r.name.Contains("Red") || r.name.Contains("Blue"))) _occludees.Add(r);

            // 给所有渲染器分配全局ID（0 留给背景）
            int nextId = 1;
            foreach (var r in _allRenderers)
                if (r) _idOf[r] = nextId++;

            Debug.Log($"All Renderers: {_allRenderers.Count}, Occludees: {_occludees.Count}");
        }

        if (GUILayout.Button("2) Bake PVS (GPU) → Save PVSData.asset"))
        {
            BakePVS_GPU();
        }

        if (GUILayout.Button("3) Attach Runtime Culler to Main Camera"))
        {
            var cam = Camera.main;
            if (!cam) { Debug.LogError("没有 MainCamera"); return; }
            var culler = cam.GetComponent<PVSRuntimeCuller>() ?? cam.gameObject.AddComponent<PVSRuntimeCuller>();
            Debug.Log("已挂 PVSRuntimeCuller。等会把 PVSData 和 Targets 指过去。");
        }
    }

    // ====== 烘焙主流程 ======
    void BakePVS_GPU()
    {
        if (_occludees.Count == 0 || _allRenderers.Count == 0) { Debug.LogError("先点 1) Collect Objects"); return; }
        if (!_cs || _kernel<0 || !_idMat) { Debug.LogError("Compute/Shader 未就绪"); return; }

        // 1) 准备 RT 与 Compute buffer
        CleanupRTAndBuffers();
        _rt = new RenderTexture(_res, _res, 24, RenderTextureFormat.ARGB32, RenderTextureReadWrite.Linear);
        _rt.enableRandomWrite = false; _rt.Create();

        _hash = new ComputeBuffer(_hashSizePow2, sizeof(uint), ComputeBufferType.Structured);
        _unique = new ComputeBuffer(_hashSizePow2, sizeof(uint), ComputeBufferType.Append);
        _unique.SetCounterValue(0);

        // 2) 暂存 & 改材质为 ObjectId
        BackupAndSwapMaterials();

        // 3) 预备数据容器
        int cells = _grid.x * _grid.y * _grid.z;
        var pvs = ScriptableObject.CreateInstance<PVSData>();
        pvs.min=_min; pvs.max=_max; pvs.grid=_grid;
        pvs.objectCount = _occludees.Count;
        for (int i=0;i<cells;i++) pvs.cellBits.Add(null);

        // 4) 相机（临时）
        var bakeCamGO = new GameObject("[PVS_BakeCamera]");
        var cam = bakeCamGO.AddComponent<Camera>();
        cam.enabled = false;
        cam.clearFlags = CameraClearFlags.SolidColor;
        cam.backgroundColor = Color.black;
        cam.cullingMask = ~0;
        cam.nearClipPlane = 0.02f;
        cam.farClipPlane = 500f;
        cam.fieldOfView = 90f;
        cam.aspect = 1f;
        cam.targetTexture = _rt;

        // 5) 循环各 cell
        var bound = new Bounds(); bound.SetMinMax(_min, _max);
        Vector3 cellSize = new(bound.size.x/_grid.x, bound.size.y/_grid.y, bound.size.z/_grid.z);
        Vector3 origin = bound.min + 0.5f * cellSize;

        // 采样方向：立方体 6 面
        Vector3[] dirs = { Vector3.forward, Vector3.back, Vector3.left, Vector3.right, Vector3.up, Vector3.down };

        int done = 0; int total = cells;
        try
        {
            for (int ix=0; ix<_grid.x; ix++)
            for (int iy=0; iy<_grid.y; iy++)
            for (int iz=0; iz<_grid.z; iz++)
            {
                int cellIdx = (ix*_grid.y + iy)*_grid.z + iz;
                Vector3 cellCenter = origin + new Vector3(ix*cellSize.x, iy*cellSize.y, iz*cellSize.z);

                // 清空 Hash 与 UniqueIDs
                ClearBuffer(_hash); _unique.SetCounterValue(0);

                // 采样点（最小实现：仅中心；你也可扩展为中心+六面+八角）
                var samplePoints = new List<Vector3>{ cellCenter };
                if (_samplesPerCell > 1)
                {
                    // 简单再加四个角点（按需要增减）
                    samplePoints.Add(cellCenter + new Vector3(+cellSize.x*0.25f, +cellSize.y*0.25f, +cellSize.z*0.25f));
                    samplePoints.Add(cellCenter + new Vector3(-cellSize.x*0.25f, +cellSize.y*0.25f, -cellSize.z*0.25f));
                    samplePoints.Add(cellCenter + new Vector3(+cellSize.x*0.25f, -cellSize.y*0.25f, -cellSize.z*0.25f));
                    samplePoints.Add(cellCenter + new Vector3(-cellSize.x*0.25f, -cellSize.y*0.25f, +cellSize.z*0.25f));
                }

                foreach (var sp in samplePoints)
                {
                    foreach (var d in dirs)
                    {
                        // 定位/朝向
                        cam.transform.position = sp;
                        cam.transform.rotation = Quaternion.LookRotation(d, Vector3.up);
                        cam.orthographic = false; // cube 面用透视 90°
                        cam.Render();             // 用当前材质（已替换成 ObjectId）渲染到 _rt

                        // 把 _rt 交给 Compute 去重
                        _cs.SetTexture(_kernel, "_IDTex", _rt);
                        _cs.SetInt("_Width", _res);
                        _cs.SetInt("_Height", _res);
                        _cs.SetInt("_HashSize", _hashSizePow2);
                        _cs.SetBuffer(_kernel, "HashTable", _hash);
                        _cs.SetBuffer(_kernel, "UniqueIDs", _unique);

                        int gx = Mathf.CeilToInt(_res/8f);
                        int gy = Mathf.CeilToInt(_res/8f);
                        _cs.Dispatch(_kernel, gx, gy, 1);
                    }
                }

                // 读回本 cell 的唯一 id 列表
                ComputeBuffer.CopyCount(_unique, _hash, 0); // 用 hash[0] 临时存数量
                uint[] countArr = new uint[1]; _hash.GetData(countArr, 0, 0, 1);
                int count = (int)countArr[0];
                uint[] ids = new uint[Mathf.Max(1,count)];
                if (count>0) _unique.GetData(ids, 0, 0, count);

                // 填位图：只关心 occludees（非 occludee 的 id 忽略）
                var bits = new bool[_occludees.Count];
                if (count>0)
                {
                    // 先做 id → occludeeIndex 的映射
                    var map = new Dictionary<uint,int>(_occludees.Count);
                    for (int i=0;i<_occludees.Count;i++)
                        if (_occludees[i]) map[(uint)_idOf[_occludees[i]]] = i;

                    for (int i=0;i<count;i++)
                    {
                        if (map.TryGetValue(ids[i], out int ocIdx)) bits[ocIdx] = true;
                    }
                }
                pvs.cellBits[cellIdx] = bits;

                done++;
                EditorUtility.DisplayProgressBar("PVS Bake (GPU)", $"Cell {done}/{total}", (float)done/total);
            }

            // 保存 asset
            string path = EditorUtility.SaveFilePanelInProject("Save PVSData", "PVSData", "asset", "");
            if (!string.IsNullOrEmpty(path))
            {
                AssetDatabase.CreateAsset(pvs, path);
                AssetDatabase.SaveAssets();
                Debug.Log($"PVSData saved: {path}");
            }
        }
        finally
        {
            EditorUtility.ClearProgressBar();
            // 还原材质并清理
            RestoreMaterials();
            if (cam) DestroyImmediate(cam.gameObject);
            CleanupRTAndBuffers();
        }
    }

    // ====== 材质替换 (SRP/Built-in 通用) ======
    void BackupAndSwapMaterials()
    {
        _backupMats.Clear();
        var mpb = new MaterialPropertyBlock();
        foreach (var r in _allRenderers)
        {
            if (!r) continue;
            _backupMats[r] = r.sharedMaterials;
            // 全部替换为同一 ID 材质（submesh 数一致）
            var arr = new Material[r.sharedMaterials.Length];
            for (int i=0;i<arr.Length;i++) arr[i] = _idMat;
            r.sharedMaterials = arr;

            // 给每个 renderer 写入自己的唯一 _ObjectId
            mpb.Clear();
            mpb.SetInt("_ObjectId", _idOf[r]);
            r.SetPropertyBlock(mpb);
        }
    }

    void RestoreMaterials()
    {
        if (_backupMats.Count==0) return;
        foreach (var kv in _backupMats)
            if (kv.Key) kv.Key.sharedMaterials = kv.Value;
        _backupMats.Clear();
    }

    // 用 SetData(0) 清空哈希表
    void ClearBuffer(ComputeBuffer buf)
    {
        // 清零
        uint[] zeros = new uint[_hashSizePow2];
        buf.SetData(zeros);
    }
}
#endif
